#include "vocab/vocab.h"
#include "model/qwen2_hub.h"

namespace pa
{

bool Vocab::engine_init_from_qwen_hub(const ModelBase& qwen2)
{
    const auto& model = (Qwen2Hub&)qwen2;

    const auto& token_strs = model.token_str();
    int token_idx = 0;
    for (token_idx = 0; token_idx < token_strs.size(); ++token_idx) {
        m_encoder[token_strs[token_idx]] = token_idx;
    }

    const auto& added_token_strs = model.added_token_str();
    for (; token_idx < (token_strs.size() + added_token_strs.size()); ++token_idx) {
        m_added_token_encoder[added_token_strs[token_idx - token_strs.size()]] = token_idx;
    }

    m_decoder = std::move(token_strs);
    m_added_token_decoder = std::move(added_token_strs);

    return true;
}

bool Vocab::encoder_init()
{
    bool flag = false;
    switch(m_model.get_type()) {
        case ModelBase::MODEL_QWEN2_HUB: {
            flag = engine_init_from_qwen_hub(m_model);
            break;
        }
        default: {
            return false;
        }
    }

    return flag;
}

const Vocab::Token& Vocab::lookup(const Id& idx) const
{
    if (idx < m_decoder.size()) {
        return m_decoder[idx];
    } else {
        return m_added_token_decoder[idx - m_decoder.size()];
    }
    PA_LOGE("idx %d is out of range", idx);
    return "";
}

Vocab::Id Vocab::lookup(const Token& token) const
{
    // unknown token?
    if (m_encoder.find(token) != m_encoder.end()) {
        return m_encoder.at(token);
    } else if (m_added_token_encoder.find(token) != m_added_token_encoder.end()) {
        return m_added_token_encoder.at(token);
    }
    return -1;
}

} // namespace pa